import numpy as np
import cv2 as cv
import os
import matplotlib.pyplot as plt
import pandas as pd


def img_norm(img):
    """
    Auth: LJH
    Args:
        img:
    Returns:
    """
    img_max = (img[img != 0]).max()
    img_min = (img[img != 0]).min()
    img_new = (img - img_min) * 65535.0 / (img_max - img_min)
    th = (0 - img_min) * 65535.0 / (img_max - img_min)
    img_new[img_new == th] = 0
    img_new = cv.medianBlur(img_new.astype(np.float32), 3)
    return img_new


def npy_data2gray(z_points, show_img=True, generate_img=False, npy_file='', target_folder=''):
    """
    Auth: WZW
    Args:
        z_points: nd.array, point cloud data
        show_img: if show image
        generate_img: if save image
        npy_file: numpy file name, used to generate image name
        target_folder: target folder to save image
    Returns:
        gray image
    """
    y_num, x_num = z_points.shape
    nan_index = np.isnan(z_points)
    # nan值灰度设置为0
    z_points[nan_index] = 0

    img_gray = img_norm(z_points)
    # 四舍五入
    # img_gray = np.around(img_gray)
    # 和三维显示对齐 需要进行旋转和翻转
    rows, cols = img_gray.shape[:2]

    # 逆时针45°旋转图片并缩小一半，第一个参数为旋转中心
    M = cv.getRotationMatrix2D((cols / 2, rows / 2), 180, 1)
    # img：源图像；M：旋转仿射矩阵；(cols,rows)：dst的大小
    img_gray = cv.warpAffine(img_gray, M, (cols, rows))

    # 翻转
    img_gray = cv.flip(img_gray, 1)
    # 转化为int8 去除小数部分
    img_gray = img_gray.astype(np.uint8)
    print(img_gray.shape)
    if show_img:
        cv.namedWindow('npy2gray', 0)
        cv.resizeWindow('npy2gray', 500, 500)
        cv.imshow('npy2gray', img_gray)
        cv.waitKey(0)
    if generate_img:
        file_name = os.path.basename(npy_file).split('.')[0] + '_gray.jpg'
        print(file_name)
        cv.imwrite(os.path.join(target_folder, file_name), img_gray)

    return img_gray


def npy_normalization_16_8(npy_data, nan_num=np.nan, imgShow=False, npySave=False, npy_file='data.npy',
                           targetFolder=""):
    """
    Auth: WZW
    PreVersion: outputNorm in V1
    Args:
        npy_data: input numpy data
        nan_num: nan is changed to a number
        imgShow: show 16 bit image and 8 bit image or not
        npySave: save normed npy file or not
        npy_file: numpy file name, used to generate saved file name
        targetFolder: path to save npy file
    Returns:
        16 bit image,
        8 bit image,
        max num in point cloud
        min num in point cloud
    """
    z = npy_data.copy()
    x = z.copy()
    nan_index = np.isnan(z)

    # 找出不是nan数值中的最大最小值
    img_max = (x[~nan_index]).max()
    img_min = (x[~nan_index]).min()

    # 归一化产生16位灰度图
    img_16bit = (z - img_min) * 65535.0 / (img_max - img_min)
    print(len(nan_index))

    # 设置nan
    if not np.isnan(nan_num):
        print("修改nan为：" + str(nan_num))
        img_16bit[nan_index] = nan_num
    # print(img_new_min)
    # print(img_new_max)
    # 产生对应的8位灰度图
    img_gray8 = img_16bit / 256
    if npySave:
        np.save(os.path.join(targetFolder, os.path.basename(npy_file).split('.')[0] + '_norm.npy'), img_16bit)
        np.save(os.path.join(targetFolder, os.path.basename(npy_file).split('.')[0] + '_norm_uint8.npy'), img_gray8)
    if imgShow:
        plt.subplot(1, 2, 1)
        plt.imshow(img_16bit, cmap='gray', vmin=0, vmax=65535)
        plt.subplot(1, 2, 2)
        plt.imshow(img_gray8, cmap='gray', vmin=0, vmax=255)
        plt.show()
    print(img_16bit.shape)
    return img_16bit, img_gray8, img_max, img_min


def npy2xls(npyData, file):
    """
    Auth: WZW
    Args:
        npyData: nd.array
        file: path to save

    Returns:

    """
    print(file)
    dataXLS = pd.DataFrame(npyData)
    writer = pd.ExcelWriter(file)
    dataXLS.to_excel(writer, 'page1')
    writer.save()


def gray2rgb_osFunction(img_gray):
    """
    Auth: WZW
    Args:
        img_gray: 8位图

    Returns:

    """
    ui8 = img_gray.astype(np.uint8)
    img_rgb = cv.cvtColor(ui8, cv.COLOR_GRAY2BGR)
    cv.namedWindow("img_rgb", 0)
    cv.resizeWindow("img_rgb", 800, 800)
    cv.imshow('img_rgb', img_rgb)
    cv.waitKey(0)
    return img_rgb


def gray2rgb_16bit(img_gray, img_show=False, generate_img=False, target_folder="", fileName=""):
    """
    Auth: WZW
    Args:
        img_gray: 灰度图
        img_show:
        generate_img:
        target_folder:
        fileName:

    Returns:

    """
    f32 = img_gray / 65535.0
    img_rgb = np.zeros((img_gray.shape[0], img_gray.shape[1], 3))
    img_rgb[:, :, 0] = f32
    img_rgb[:, :, 1] = f32
    img_rgb[:, :, 2] = f32
    if img_show:
        # plt.imshow(img_rgb)
        # plt.show()
        cv.namedWindow("16bit", 0)
        cv.resizeWindow('16bit', 800, 800)
        cv.imshow('16bit', img_rgb)
        cv.waitKey(0)
    if generate_img:
        img_rgb *= 255.0
        img_rgb = img_rgb.astype(np.uint8)
        cv.imwrite(os.path.join(target_folder, fileName), img_rgb)
    return img_rgb


def gray2rgb_8bit(img_gray, img_show=False, generate_img=False, target_folder="", fileName=""):
    f32 = img_gray.copy()

    # 拷贝rgb通道，效果差
    # img_rgb = np.zeros((img_gray.shape[0], img_gray.shape[1], 3))
    # img_rgb[:, :, 0] = f32
    # img_rgb[:, :, 1] = f32
    # img_rgb[:, :, 2] = f32

    # 系统转换
    ui8 = f32.astype(np.uint8)
    img_rgb = cv.cvtColor(ui8, cv.COLOR_GRAY2BGR)
    if img_show:
        # plt.imshow(img_rgb)
        # plt.show()
        cv.namedWindow("8bit", 0)
        cv.resizeWindow('8bit', 800, 800)
        cv.imshow('8bit', img_rgb)
        cv.waitKey(0)
    if generate_img:
        img_rgb = img_rgb.astype(np.uint8)
        cv.imwrite(os.path.join(target_folder, fileName), img_rgb)
    return img_rgb


def astypeTest():
    """
    used to detect the function "astype"
    """
    da = np.array([0, 1, 2, 256, 257, 10000, 20000, 30000, 40000, 50000, 60000, 65535, 65536])
    da16 = da.astype(np.uint16)
    da8 = da.astype(np.uint8)
    print(da)
    print(da16)
    print(da8)


def _get_pixel(event, x, y, flags, param):
    global img
    # 画笔相关
    circleRadius = 5
    point_color = (255, 255, 0)
    thickness = 2
    point_start = (0, 0)
    point_end = (0, 0)
    if event == cv.EVENT_LBUTTONDOWN:
        # 检测到鼠标左键按下，按下就画圆，并且记录抬起按键之前的点
        point_start = (x, y)
        cv.circle(img, point_start, int(circleRadius / 2), point_color, -1)
    elif event == cv.EVENT_LBUTTONUP:
        # 检测到鼠标左键抬起
        point_end = (x, y)
        print("point_end is {}".format(point_end))


def get_pixel(img_path):
    """
    通过可视化点击，获取图像上的像素位置
    :param img_path:
    :return:
    """
    global img
    img = cv.imread(img_path)
    cv.namedWindow('get_pixel', cv.WINDOW_NORMAL)
    cv.setMouseCallback('get_pixel', _get_pixel)
    while True:
        cv.imshow('get_pixel', img)
        if (cv.waitKey(1) == ord('q')):
            break
    cv.destroyAllWindows()


if __name__ == "__main__":
    data = np.load(
        "D:\\Programs\\data\\dataset\\20220628-full-size\\2022-06-23\\Up\\35051772623505277134A23062200150556_up_2022-06-23-15-39-06.npy")

    npy_normalization_16_8(npy_data=data, nan_num=np.nan, imgShow=True)
